﻿using Ndns.Packet.Abstractions;
using System;
using System.Collections.Generic;
using System.IO;

namespace Ndns.Packet.Parser
{
    public static class DnsPacketExtensions
    {
        public static DnsPacket CloneAsResponse(this DnsPacket request, ushort transactionID, DnsReplyCode replyCode)
        {
            var packet = new DnsPacket(transactionID, request.Flag);
            packet.IsResponsePacket = true;
            packet.ReplyCode = replyCode;
            foreach (var item in request.Queries) packet.Queries.Add(new DnsQuery() { HostName = item.HostName, Type = item.Type, Class = item.Class });
            return packet;
        }
        public static DnsPacket CloneAsResponse(this DnsPacket request, DnsReplyCode replyCode)
        {
            return CloneAsResponse(request, request.TransactionID, replyCode);
        }

        public static int CopyTo(this DnsPacket packet, Stream stream)
        {
            int count;
            int offset;

            //DNS数据包头
            stream.WriteNetworkUInt16(packet.TransactionID);
            stream.WriteNetworkUInt16((ushort)packet.Flag);

            //各个记录的数量
            stream.WriteNetworkUInt16((ushort)packet.Queries.Count);
            stream.WriteNetworkUInt16((ushort)packet.AnswerRRs.Count);
            stream.WriteNetworkUInt16((ushort)packet.AuthoritativeNSs.Count);
            stream.WriteNetworkUInt16((ushort)packet.AdditionalRRs.Count);

            count = 12;
            offset = 12;

            //各条记录
            foreach (var e in packet.Queries)
                count += e.CopyTo(stream, ref offset);
            foreach (var e in packet.AnswerRRs)
                count += e.CopyTo(stream, ref offset);
            foreach (var e in packet.AuthoritativeNSs)
                count += e.CopyTo(stream, ref offset);
            foreach (var e in packet.AdditionalRRs)
                count += e.CopyTo(stream, ref offset);

            return count;
        }

        public static int CopyCompressedTo(this DnsPacket packet, Stream stream)
        {
            int count;
            int offset;
            var symbols = new Dictionary<string, DnsHostNamePointer>();

            //DNS数据包头
            stream.WriteNetworkUInt16(packet.TransactionID);
            stream.WriteNetworkUInt16((ushort)packet.Flag);

            //各个记录的数量
            stream.WriteNetworkUInt16((ushort)packet.Queries.Count);
            stream.WriteNetworkUInt16((ushort)packet.AnswerRRs.Count);
            stream.WriteNetworkUInt16((ushort)packet.AuthoritativeNSs.Count);
            stream.WriteNetworkUInt16((ushort)packet.AdditionalRRs.Count);

            count = 12;
            offset = 12;

            //各条记录
            foreach (var e in packet.Queries)
                count += e.CopyCompressedTo(stream, symbols, ref offset);
            foreach (var e in packet.AnswerRRs)
                count += e.CopyCompressedTo(stream, symbols, ref offset);
            foreach (var e in packet.AuthoritativeNSs)
                count += e.CopyCompressedTo(stream, symbols, ref offset);
            foreach (var e in packet.AdditionalRRs)
                count += e.CopyCompressedTo(stream, symbols, ref offset);

            return count;
        }


        public static byte[] GetBytes(this DnsPacket packet, bool allowTruncation = false, bool enableCompression = true)
        {
            byte[] buff;
            //获得字节
            using (var ms = new MemoryStream(512))
            {
                if (enableCompression)
                    packet.CopyCompressedTo(ms);
                else
                    packet.CopyTo(ms);
                buff = ms.ToArray();
            }
            //处理截断
            if (allowTruncation)
            {
                if (buff.Length > 512)
                {
                    var truncatedBuff = new byte[512];
                    Buffer.BlockCopy(buff, 0, truncatedBuff, 0, 512);
                    buff = truncatedBuff;
                }
            }
            return buff;
        }
    }
}
